# -*- coding: utf-8 -*-
# @Time    : 2020/6/14 20:12
# @Author  : Zhang Haozhen
# @FileName: light_recog.py

import numpy as np
import cv2
import collections
import time
import datetime

# 生成HSV色域范围字典
def getColorList(sat_low, val_low):
    '''
    :return: 红绿黄三色色域字典
    '''
    dict = collections.defaultdict(list)

    # 红色
    lower_red = np.array([0, sat_low, val_low])
    upper_red = np.array([10, 255, 255])
    color_list = []
    color_list.append(lower_red)
    color_list.append(upper_red)
    lower_red = np.array([150, sat_low, val_low])
    upper_red = np.array([180, 255, 255])
    color_list.append(lower_red)
    color_list.append(upper_red)
    dict['Red'] = color_list

    # 黄色
    lower_yellow = np.array([26, sat_low, val_low])
    upper_yellow = np.array([40, 255, 255])
    color_list = []
    color_list.append(lower_yellow)
    color_list.append(upper_yellow)
    lower_yellow = np.array([11, sat_low, val_low])
    upper_yellow = np.array([25, 255, 255])
    color_list.append(lower_yellow)
    color_list.append(upper_yellow)
    dict['Yellow'] = color_list

    # 绿色
    lower_green = np.array([35, sat_low, val_low])
    upper_green = np.array([90, 255, 255])
    color_list = []
    color_list.append(lower_green)
    color_list.append(upper_green)
    dict['Green'] = color_list

    return dict

# 判断颜色
def get_color(frame, color_list):
    '''
    :param frame: 待检测图像
    :return: 颜色
    '''
    if frame is None:
        return 'Red'
    hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
    
    # The following lines are moved to main_traffic_light_recognition.py for dynamic adjustment
    # sum_saturation = np.sum(hsv[:,:,1])
    # area = frame.shape[0] * frame.shape[1]
    # avg_saturation = sum_saturation / area
    # sat_low = int(avg_saturation*1.3)#均值的1.3倍，工程经验
    # val_low = 140


    maxsum = -100
    color = None
    # color_dict = getColorList(sat_low)
    color_dict = color_list # Now getColorList is called in main_traffic_light_recognition.py
    for d in color_dict:
        sum = 0
        mask = None
        for i in range(len(color_dict[d]) // 2):
            if mask is not None:
                mask += cv2.inRange(hsv, color_dict[d][2 * i], color_dict[d][2 * i + 1])
            else:
                mask = cv2.inRange(hsv, color_dict[d][2 * i], color_dict[d][2 * i + 1])
        binary = cv2.threshold(mask, 127, 255, cv2.THRESH_BINARY)[1]
        binary = cv2.dilate(binary, None, iterations=2)
        cnts, hiera = cv2.findContours(binary.copy(), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        for c in cnts:
            sum += cv2.contourArea(c)

        if sum > maxsum:
            maxsum = sum
            color = d

    return color


if __name__ == '__main__':
    # H（色调）、S（饱和度）、V（明度）
    frame = cv2.imread("./sample/yellow4.png")
    prev = time.time()
    print(get_color(frame))
    after = time.time()
    inference_time = datetime.timedelta(seconds=after - prev)
    print(inference_time)